import importlib

import numpy as np
import torch
import torch.nn.functional as F
import torch.optim as optim

from defense.base_defense import BaseDefense
from raw.base_model import BaseModel

class Defense(BaseDefense):
    def __init__(self, model_name, config):
        super().__init__('EAT', model_name, config, retrain = True)

    def get_external_model_group(self):
        name = self.model.model_name
        models = importlib.import_module('defense.methods.EAT_external_models.' + name).get_model_group()
        for model in models:
            assert isinstance(model, BaseModel), 'Model class must extend raw.base_model.BaseModel'
            assert model.model_name == name, 'Model name must be the same as its file name'
            model.to_device()
        return models

    def train_external_model_group(self, train_loader, valid_loader):
        models = self.get_external_model_group()

        for i in range(len(models)):
            opt = models[i].get_optimizer()

            print('Training external model #{}:\n'.format(i))

            best_ratio = 0
            best_state_dict = None
            for epoch in range(self.model.training_params['num_epochs']):
                models[i].train_once(train_loader, opt, epoch)
                ratio = models[i].test(valid_loader)
                models[i].adjust_optimizer(opt, epoch)
                if ratio > best_ratio:
                    best_ratio = ratio
                    best_state_dict = models[i].state_dict_clone()
                    print('Validate: Accuracy increased')
                else:
                    print('Validate: Accuracy decreased')
                print()

            path = 'data/enhanced_models/{}_{}_{}.pt'.format(self.model.model_name, self.name, i)
            torch.save(best_state_dict, path)
            print('Model saved to ' + path + '\n')

    def load_external_model_group(self, test_loader):
        models = self.get_external_model_group()

        for i in range(len(models)):
            models[i].load_state_dict(torch.load('data/enhanced_models/{}_{}_{}.pt'.format(self.model.model_name,
                                                                                            self.name, i)))
            models[i].test(test_loader)
        print()

        return models

    # 进行RFGSM攻击
    def RFGSM_attack(self, model, images):
        model.eval()

        with torch.no_grad():
            random_sign = torch.sign(torch.randn(*images.size())).to(self.model.device)
            new_images = torch.clamp(images + self.config['alpha'] * random_sign, min = 0.0, max = 1.0)

        new_images.requires_grad = True

        logits_attack = model(new_images)
        # To avoid label leaking, we use the model's output instead of the true labels
        labels_attack = torch.argmax(logits_attack, dim = 1)
        loss_attack = F.cross_entropy(logits_attack, labels_attack)
        gradient = torch.autograd.grad(loss_attack, new_images)[0]

        new_images.requires_grad = False

        # generation of adversarial examples
        with torch.no_grad():
            xs_adv = new_images + (self.config['eps'] - self.config['alpha']) * torch.sign(gradient)
            xs_adv = torch.clamp(xs_adv, min = 0.0, max = 1.0)
        return xs_adv

    # 单轮防御训练
    def train(self, pre_trained_models, train_loader, opt, epoch):
        for ind, (images, labels) in enumerate(train_loader):
            images = images.to(self.model.device)
            labels = labels.to(self.model.device)

            # 随机选择要被攻击的模型
            i = np.random.randint(len(pre_trained_models) + 1)
            attacking_model = self.model if i == 0 else pre_trained_models[i - 1]

            # 进行RFGSM攻击
            advs = self.RFGSM_attack(attacking_model, images)

            self.model.train()

            output = self.model(images)
            loss = F.cross_entropy(output, labels) #原始图像损失

            adv_output = self.model(advs)
            adv_loss = F.cross_entropy(adv_output, labels) #混淆图像损失

            total_loss = 0.5 * (loss + adv_loss) #总损失

            opt.zero_grad()
            total_loss.backward()
            opt.step()

            ind += 1
            print('\rEpoch {}: [batch: {}/{} ({:.0f}%)] \tLoss: {:.4f} AdvLoss: {:.4f} TotalLoss: {:.4f}'.format(
                epoch, ind, len(train_loader),
                ind / len(train_loader) * 100, loss, adv_loss, total_loss), end=' ')

        print()

    def defend(self):
        batch_size = self.model.training_params['batch_size']
        # 训练数据
        train_loader, valid_loader = self.model.get_train_valid_loader(batch_size = batch_size,
                                                                        valid_size = 0.1,
                                                                        transform_train = True,
                                                                        shuffle = True)
        # 测试数据
        test_loader = self.model.get_test_loader(batch_size = batch_size, shuffle = False)
        # 获取optimizer
        opt = self.model.get_optimizer()

        # 训练外部模型
        if self.config['train_externals']:
            self.train_external_model_group(train_loader, valid_loader)

        # 加载外部模型
        models = self.load_external_model_group(test_loader)

        # 防御训练
        best_ratio = 0
        best_state_dict = None
        for epoch in range(self.model.training_params['num_epochs']):
            self.train(models, train_loader, opt, epoch)
            ratio = self.model.test(valid_loader)
            self.model.adjust_optimizer(opt, epoch)
            if ratio > best_ratio:
                best_ratio = ratio
                best_state_dict = self.model.state_dict_clone()
                print('Validate: Accuracy increased')
            else:
                print('Validate: Accuracy decreased')
            print()
        return best_state_dict
